
from twisted.protocols.basic import Int16StringReceiver
from twisted.internet.protocol import Factory
from twisted.internet import reactor

from twisted.python import log

import hashlib

import time
import srp
import db
import sqlite3

from google.protobuf import message

from srp_auth import srp_auth_pb2
from srp_auth import ocrypt



class InvalidMessage (Exception):
    pass

class AuthenticationFailed (Exception):
    pass

class AccessDenied (Exception):
    pass

class RequestError (Exception):
    pass


def assert_fields(msg, required_fields):
    for f in required_fields.split():
        if f:
            if not msg.HasField(f):
                raise InvalidMessage

def phex(s):
    return ''.join( ['%x' % ord(c) for c in s] )


            
class SrpAuthServerProtocol (Int16StringReceiver):
                
    def connectionMade(self):
        self.auth_step     = 1
        self.authenticated = False
        self.username      = None
        self.is_admin      = None
        self.is_locked     = None
        self.expires_at    = None
        self.srp_key       = None
        
        self.handlers = { srp_auth_pb2.ADD_USER            : self.add_user,
                          srp_auth_pb2.REMOVE_USER         : self.remove_user,
                          srp_auth_pb2.GET_USER_LIST       : self.get_users,
                          srp_auth_pb2.GET_USER_PROPERTIES : self.get_user_props,
                          srp_auth_pb2.SET_USER_PROPERTIES : self.set_user_props,
                          srp_auth_pb2.CHANGE_PASSWORD     : self.change_password,
                          srp_auth_pb2.GET_PUBLIC_INFO      : self.get_public_info }
                          

    def sendErrorMessage(self, msg, drop_connection=True):
        reply               = srp_auth_pb2.SRPAuthMessage()
        reply.msg_type      = srp_auth_pb2.REPLY
        reply.error_message = msg
        
        self.sendString( reply.SerializeToString() )

        if drop_connection:
            self.transport.loseConnection()
            #print 'Dropping connection due to: ', msg
      
      
    def stringReceived(self, s):
        try:
            msg = srp_auth_pb2.SRPAuthMessage()
            msg.ParseFromString( s )
            
            reply = srp_auth_pb2.SRPAuthMessage()
            
            if not self.authenticated:
                if (msg.msg_type == srp_auth_pb2.AUTHENTICATION and msg.HasField('auth')):
                    self.handleAuthMessage( msg.auth, reply )
                    
                elif (msg.msg_type == srp_auth_pb2.GET_PUBLIC_INFO):
                    reply.msg_type = srp_auth_pb2.REPLY
                    self.get_public_info( msg, reply )
                    self.sendString( reply.SerializeToString() )
                    self.transport.loseConnection()
                    
                else:
                    raise AuthenticationFailed
            else:

                handler = self.handlers.get( msg.msg_type, None )

                if not handler:
                    raise InvalidMessage
                
                reply.msg_type = srp_auth_pb2.REPLY
                handler( msg, reply )
            
            self.sendString( reply.SerializeToString() )
            
        except InvalidMessage:
            self.sendErrorMessage('Invalid protobuf message content')
        except message.Error:
            self.sendErrorMessage('Invalid protobuf message')
        except AuthenticationFailed:
            self.sendErrorMessage('Authentication failed')
            log.msg('Authentication failed for user %s' % str(self.username))
        except AccessDenied:
            self.sendErrorMessage('Access Denied', False)
            log.msg('Privileged operation denied for user %s' % str(self.username))
        except RequestError, e:
            self.sendErrorMessage(str(e), False)
        except sqlite3.Error, e:
            self.sendErrorMessage('Internal database error: ' + str(e), False)
            log.msg('Internal database error: ' + str(e))
            
         
    def handleAuthMessage(self, msg, reply):
        reply.msg_type = srp_auth_pb2.AUTHENTICATION
    
        tocli = reply.auth
         
        if self.auth_step == 1:
         
            assert_fields( msg, 'username A' )

            tpl = self.factory.db.get_user( msg.username )
        
            if tpl[0] is None:
                raise AuthenticationFailed
        
            self.username      = msg.username
            
            salt               = tpl[0]
            verification_key   = tpl[1]
            
            self.is_locked     = tpl[2]
            self.expires_at    = tpl[3]
            self.is_admin      = tpl[4]
            
            #print 'Username: ', self.username, 'IS_ADMIN: ', self.is_admin
            
            #salt, verification_key = srp.create_salted_verification_key( 'thomas.cocagne', 'foo' )

            #print 'Authkey: ', phex(verification_key)
            
            self.svr  = srp.Verifier( self.username, salt, verification_key, msg.A,
                                      hash_alg = self.factory.db.srp_hash_type,
                                      ng_type  = self.factory.db.srp_prime_type )   
        
            s, B = self.svr.get_challenge()
            
            if s is None:
                raise AuthenticationFailed
            
            tocli.s = s
            tocli.B = B
        
            self.auth_step = 2
         
        elif self.auth_step == 2:
            
            self.auth_step = -1

            assert_fields( msg, 'M' )
            
            HAMK = self.svr.verify_session( msg.M )
            
            if HAMK is None:
                raise AuthenticationFailed
            
            tocli.HAMK        = HAMK
            reply.signed_HAMK = self.factory.pri_key.sign_oneshot( HAMK )

            # negative expire time means infinite
            if self.expires_at >= 0:
                tocli.password_expired = time.time() > self.expires_at


            log.msg('Authenticated user %s' % str(self.username))
            
            self.authenticated = True
            self.srp_key       = self.svr.get_session_key()
            
        else:
            raise AuthenticationFailed
      
      
    def add_user(self, msg, reply):
        if not self.is_admin:
            raise AccessDenied
        assert_fields( msg, 'add_user' )
      
        au = msg.add_user
        up = msg.add_user.properties
      
        username   = up.username
        salt       = au.salt
        verifier   = ocrypt.aes_oneshot_decrypt(self.srp_key, salt, au.verification_key)
        is_locked  = False
        expires_at = 0
        is_admin   = False
        
        if up.HasField('is_locked'):
            is_locked = up.is_locked
            
        if up.HasField('expires_at'):
            expires_at = up.expires_at
         
        if up.HasField('is_admin'):
            is_admin = up.is_admin

        try:
            self.factory.db.add_user( username, salt, verifier, is_locked, expires_at, is_admin )
            log.msg('New user %s added by %s' % (username, self.username))
        except db.DBError, e:
            raise RequestError(str(e))
   
      
    def remove_user(self, msg, reply):
        if not self.is_admin:
            raise AccessDenied
        assert_fields( msg, 'username' )
      
        try:
            self.factory.db.delete_user( msg.username )
            log.msg('User %s deleted by %s' % (msg.username, self.username))
        except db.DBError, e:
            raise RequestError(str(e))

        
    def get_users(self, msg, reply):
        if not self.is_admin:
            raise AccessDenied
        for uname in self.factory.db.get_user_list():
            reply.user_list.append( str(uname) )

   
    def get_user_props(self, msg, reply):
        if not self.is_admin:
            raise AccessDenied
        assert_fields( msg, 'username' )
        
        salt, verifier, is_locked, expires_at, is_admin = self.factory.db.get_user( msg.username )
        
        if salt is None:
            RequestError("Invalid username: " + msg.username)

        up = reply.user_properties
        
        up.username   = msg.username
        up.is_locked  = is_locked
        up.expires_at = expires_at
        up.is_admin   = is_admin
        
        
    def set_user_props(self, msg, reply):
        if not self.is_admin:
            raise AccessDenied
        assert_fields( msg, 'user_properties' )
        
        up = msg.user_properties
        
        d = self.factory.db
        
        if up.HasField('is_locked'):
            d.set_user_is_locked( up.username, up.is_locked )
            log.msg( '%s %s account %s' % (self.username, 'locked' if up.is_locked else 'unlocked', up.username))
         
        if up.HasField('expires_at'):
            xp = int(time.time() + up.expires_at)
            d.set_user_expires_at( up.username, xp )
            log.msg( '%s changed expiry of account %s to %d' % (self.username, up.username, xp))
            
        if up.HasField('is_admin'):
            d.set_user_is_admin( up.username, up.is_admin )
            log.msg('%s %s administrative privileges for %s' %
                    (self.username, 'granted' if up.is_admin else 'revoked', up.username))
            
            

    def change_password(self, msg, reply):
        assert_fields( msg, 'change_password' )
        cp = msg.change_password
        
        username = self.username

        if cp.HasField('username'):
            if not cp.username == self.username and not self.is_admin:
                raise AccessDenied
            username = cp.username

        new_vk = ocrypt.aes_oneshot_decrypt( self.srp_key, cp.salt, cp.verification_key )
         
        self.factory.db.change_user_verifier( username, cp.salt, new_vk )

        log.msg('Password changed for %s by %s' % (username, self.username))


    def get_public_info(self, msg, reply):
        reply.public_info.public_key_der = self.factory.db.pub_key_bytes
        reply.public_info.hash_alg       = self.factory.db.srp_hash_type
        reply.public_info.prime_size     = self.factory.db.srp_prime_type
      
         


            
class SrpAuthFactory ( Factory ):
            
    protocol = SrpAuthServerProtocol
    
    def __init__(self, db_filename):
        self.db = db.DB( db_filename )
        self.pri_key = ocrypt.PrivateKey( self.db.pri_key_bytes )
        
   
